# -*- coding: utf-8 -*-
"""
Created on Wed Apr 14 12:31:12 2021
# Description :构建受限玻尔兹曼机
@author: Lenovo-pc
"""

import sys
import os

import numpy as np
np.random.seed(1)

class rbm_binary:
    def __init__(self, num_visible, num_hidden, learning_rate = 0.1):
        self.num_visible = num_visible # 可见层维度 , 下标i,v行
        self.num_hidden = num_hidden # 隐藏层维度， 下标j，h列
        self.learning_rate = learning_rate # 学习率

        # 权重矩阵W，维度（num_visible*num_hidden）,初始值使用高斯分布随机生成（均值0，方差0.1）
        self.weights = 0.1*np.random.randn(self.num_visible,self.num_hidden) #v*h
        self.weights = np.insert(self.weights,0,0,axis=0) # 在W的第一行增加可见层的偏置a ，(v+1)*h
        self.weights = np.insert(self.weights,0,0,axis=1) # 在W的第一列增加隐藏层的偏置b ，(v+1)*(h+1)

    def train(self
              ,datasets=np.zeros([2,3])  # d*v,d:样本量
              ,batch_size=1
              ,max_epoch=1
              ):
        """
        Train the machine
        :param datasets: 样本数据集，每一行是一个可见层示例
        :param batch_size:小批量数据大小
        :param max_epoch:所有样本全部学习次数

        """
        num_examples = datasets.shape[0]
        num_batches = num_examples//batch_size
        
        data = np.insert(datasets,0,1,axis=1) #插入偏置，初始1，d*(v+1)
       

        for epoch in range(max_epoch):
            error = 0
            for K in range(num_batches):
                batch_data = data[K*batch_size:(K+1)*batch_size]
                d_weights,d_error = self.CDK(batch_size,batch_data)
                self.weights += self.learning_rate*d_weights
                error += d_error
            #print("Epoch %s: error is %s" % (epoch, error))

    #cd-k散度
    def CDK(self
            ,batch_size
            ,datasets=np.zeros([2,3])  # batch*(v+1)
            ,k=1):

        def hidden_prob(visible=np.array([0])):
            return  self._logistic(np.dot(visible, self.weights))# d*(h+1)

        visible= datasets  # d*(v+1)
        pos_hidden_prob = hidden_prob(visible)

        for t in range(k): # k次状态转移
            hidden = hidden_prob(visible) > np.random.rand(batch_size,self.num_hidden+1)
            hidden = hidden+0


            visible_prob = self._logistic(np.dot(hidden,self.weights.T) ) #d*(v+1)
            visible = visible_prob>np.random.rand(batch_size,self.num_visible+1)
            #visible = visible_prob
            visible[:, 0] = 1  # 保证可见层偏置a可用

        neg_hidden_prob = hidden_prob(visible)
        d_weights = (np.dot(datasets.T,pos_hidden_prob)-np.dot(visible.T,neg_hidden_prob))/batch_size
        d_error = np.sum((datasets - visible) ** 2)
        return d_weights,d_error

    def get_hidden(self,data):
        """
        :param data: 可见层
        :return: 隐藏层
        """
        num_examples = data.shape[0]
        hidden_states = np.ones((num_examples, self.num_hidden + 1))
        data = np.insert(data, 0, 1, axis = 1)
        hidden_activations = np.dot(data, self.weights)
        hidden_probs = self._logistic(hidden_activations)
        hidden_states[:,:] = hidden_probs > np.random.rand(num_examples, self.num_hidden + 1)

        # 去除偏置位
        hidden_states = hidden_states[:,1:]
        return hidden_states


    def get_visible(self,data):
        num_examples = data.shape[0]
        visible_states = np.ones((num_examples, self.num_visible + 1))
        data = np.insert(data, 0, 1, axis = 1)
        visible_activations = np.dot(data, self.weights.T)
        visible_probs = self._logistic(visible_activations)
        visible_states[:,:] = visible_probs > np.random.rand(num_examples, self.num_visible + 1)

        # 去除偏置位
        visible_states = visible_states[:,1:]
        return visible_states

    def _logistic(self, x):
        return 1.0 / (1 + np.exp(-x))

if __name__ == '__main__':
    r = rbm_binary(num_visible = 6, num_hidden = 6)
    training_data = np.array([[1,1,1,0,0,0],
                              [1,0,1,0,0,0],
                              [1,1,1,0,0,0],
                              [0,0,1,1,1,0],
                              [0,0,1,1,0,0],
                              [0,0,1,1,1,0]])
    r.train(training_data,batch_size=2,max_epoch=5000)
    print(r.weights)

    user = np.array([[0,0,0,1,1,0]])
    print(r.get_hidden(user))
    # print('========Generating Samples========')
    # print(r.get_visible(r.get_hidden(user)))
